#include <cstdio>
#include <cstdlib>
#include <pthread.h>
#include <mqueue.h>

#define ARRAY_SIZE 10000  // 数组大小
#define THREAD_COUNT 4    // 线程数量

int *arr;
int *squaredArr;  // 存储平方结果的数组

// 消息队列属性
struct mq_attr attr = {
        .mq_flags = 0,  // 消息队列标志
        .mq_maxmsg = THREAD_COUNT, // 最大消息数
        .mq_msgsize = sizeof(int),  // 消息大小
        .mq_curmsgs = 0  // 当前队列中的消息数
};


// 消息队列
mqd_t mq;

// 线程执行函数：计算数组的一部分的平方
void *calculate_square(void *arg)
{
    int thread_id = *((int *) arg);  // 获取线程ID
    int start = (ARRAY_SIZE / THREAD_COUNT) * thread_id;
    int end = (thread_id == THREAD_COUNT - 1) ? ARRAY_SIZE : (ARRAY_SIZE / THREAD_COUNT) * (thread_id + 1);

    for (int i = start; i < end; i++)
    {
        squaredArr[i] = arr[i] * arr[i];  // 计算平方并存储
    }

    // 发送消息到消息队列，表示当前子线程已完成工作
    mq_send(mq, (char *) &thread_id, sizeof(thread_id), 0);

    return NULL;
}

int test01()
{
    // 动态分配数组内存
    arr = (int *) malloc(ARRAY_SIZE * sizeof(int));
    squaredArr = (int *) malloc(ARRAY_SIZE * sizeof(int));

    if (arr == NULL || squaredArr == NULL)
    {
        perror("Failed to allocate memory for arrays");
        exit(1);
    }

    // 填充数组，数组内容为从 1 到 10000 的整数
    for (int i = 0; i < ARRAY_SIZE; i++)
    {
        arr[i] = i + 1;  // 数组元素从 1 开始
    }

    // 创建消息队列
    mq = mq_open("/my_queue", O_CREAT | O_RDWR, 0644, &attr);
    if (mq == (mqd_t) -1)
    {
        perror("Failed to open message queue");
        exit(1);
    }

    // 创建线程数组
    pthread_t threads[THREAD_COUNT];
    int thread_ids[THREAD_COUNT];

    // 启动线程，分配任务
    for (int i = 0; i < THREAD_COUNT; i++)
    {
        thread_ids[i] = i;
        if (pthread_create(&threads[i], NULL, calculate_square, (void *) &thread_ids[i]) != 0)
        {
            perror("Failed to create thread");
            exit(1);
        }
    }

    // 等待所有线程结束
    int msg_buffer;
    for (int i = 0; i < THREAD_COUNT; i++)
    {
        // 接收消息队列中的消息
        if (mq_receive(mq, (char *) &msg_buffer, sizeof(msg_buffer), NULL) == -1)
        {
            perror("Failed to receive message");
            exit(1);
        }
    }

    // 所有子线程已完成工作，可以输出结果
    printf("First 10 squared elements of the array:\n");
    for (int i = 0; i < 10; i++)
    {
        printf("%d ", squaredArr[i]);
    }
    printf("\n");

    // 释放内存
    free(arr);
    free(squaredArr);

    // 关闭并删除消息队列
    mq_close(mq);
    mq_unlink("/my_queue");

    return 0;
}
